package io.youngledo.usb.task;

import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.function.Consumer;

/**
 * 支持进度跟踪的输出流包装器
 * 在写入数据时自动更新进度跟踪器
 */
public class ProgressOutputStream extends FilterOutputStream {

    private final ProgressUpdater progressUpdater;

    /**
     * 使用ProgressTracker构造
     */
    public ProgressOutputStream(OutputStream outputStream, ProgressTracker progressTracker) {
        super(outputStream);
        this.progressUpdater = new ProgressUpdater(progressTracker, () -> {
            try {
                close();
            } catch (IOException e) {
                // 忽略关闭时的异常
            }
        });
    }

    /**
     * 使用回调函数构造
     */
    public ProgressOutputStream(OutputStream outputStream, Consumer<Long> progressCallback) {
        super(outputStream);
        this.progressUpdater = new ProgressUpdater(progressCallback, null);
    }

    @Override
    public void write(int b) throws IOException {
        super.write(b);
        progressUpdater.updateProgress(1);
    }

    @Override
    public void write(byte[] buffer) throws IOException {
        write(buffer, 0, buffer.length);
    }

    @Override
    public void write(byte[] buffer, int offset, int length) throws IOException {
        super.write(buffer, offset, length);
        progressUpdater.updateProgress(length);
    }

    @Override
    public void flush() throws IOException {
        super.flush();
        // 刷新时不调用finalizeProgress，避免重复计算
        // finalizeProgress只在close时调用
    }

    @Override
    public void close() throws IOException {
        try {
            // 刷新缓冲区
            flush();
            // 最终进度更新在flush中已完成
        } finally {
            super.close();
        }
    }

    public long getTotalBytesWritten() {
        return progressUpdater.getTotalBytesProcessed();
    }
}